#!/home/zhonc0a/Enthought/Canopy_64bit/User/bin/python
import sys
import re
import matplotlib
matplotlib.use('QT4Agg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
import numpy as np
import argparse
        
# locate the position of necessary column (mol type mass xu yu zu ix iy iz)

    
    
class massCenterGroup():
    def __init__(self,inputCommand):
        self.groupName=inputCommand.split(':')[0]
        self.groupType=inputCommand.split(':')[1]
        self.atomType=inputCommand.split(':')[2].split(',')
        if self.groupType.upper()=='C':
            try:
                self.clusterRange=float(inputCommand.split(':')[3])
            except IndexError:
                print "for cluster type, the range should be input. Correct format NAME:GROUPTYPE:ATOMTYPE:RANGE:OTHERATOMTYPE"
                sys.exit()
        try:
            self.atomTypeEx=inputCommand.split(':')[4]
        except IndexError:
            self.atomTypeEx=[]
       #self.polyConnectDirection=inputCommand.split(':')[5].split(',')[0]
       #self.polyConnectAtom=inputCommand.split(':')[5].split(',')[1:]
        self.molID=[]
        self.massCenter={}
        self.atomList=[]
        self.atomListEx=[]
        self.molList={}
        self.clusterList={}
        
       
    def extractInfo(self,trjData,indexList,characterLen):
        cellParam=[]
        for line in trjData:
            if len(line.split())==characterLen and line.split()[indexList[1]] in self.atomType:
                self.atomList.append([line.split()[i] for i in indexList])  
            if len(line.split())==characterLen and line.split()[indexList[1]] in self.atomTypeEx:
                self.atomListEx.append([line.split()[i] for i in indexList])    
            if len(line.split())==2 and re.search('^-?\d+\.\d+\s+-?\d+\.\d+\s*$',line):
                cellParam.append(map(float,line.split()))
        #molList is a dictionary that use mol id as index and selected atom as value
        
        
        for line in self.atomList:
            if line[0] not in self.molList:
                self.molList[line[0]]=[map(float,line)]
            else:
                self.molList[line[0]].append(map(float,line))
        #if group type is molecule, remove molecule that doesn't have all the listed type.
        deleteList=[]
        if self.groupType.upper()=='M':
            for mol in self.molList:
                typeList=[]
                ifDelete=0
                for atom in self.molList[mol]:
                    typeList.append(atom[1])
                for atomType in self.atomType:
                    if float(atomType) not in typeList:
                        ifDelete=1
                if ifDelete==1:
                    deleteList.append(mol)
            for mol in deleteList:
                self.molList.pop(mol)
            if not self.molList:
                print 'Not molecule can be indentified as '+self.groupName+', exit now'
                sys.exit()                    
            #print 'Molecules that identified as '+self.groupName+' are '+','.join([mol for mol in self.molList])+'. And the number of atoms in them are '+','.join([str(len(self.molList[mol])) for mol in self.molList])             
 
                
        #calculate cell length
        cellX=cellParam[0][1]-cellParam[0][0]
        cellY=cellParam[1][1]-cellParam[1][0]
        cellZ=cellParam[2][1]-cellParam[2][0]
        self.cellX=cellX
        self.cellY=cellY
        self.cellZ=cellZ
        ##shift the whole group if any element in group is out of bound
        ##first determine the number of shift cell in each direction
        shiftMarkerList={}
        shiftMarker={}
        for mol in self.molList:
            shiftMarkerList[mol]=[[],[],[]]
            for atom in self.molList[mol]:
                shiftMarkerList[mol][0].append(atom[6])
                shiftMarkerList[mol][1].append(atom[7])
                shiftMarkerList[mol][2].append(atom[8])
        for mol in shiftMarkerList:    
            shiftX=min(shiftMarkerList[mol][0])
            if abs(min(shiftMarkerList[mol][0]))<abs(max(shiftMarkerList[mol][0])):
                shiftX=max(shiftMarkerList[mol][0])
            shiftY=min(shiftMarkerList[mol][1])
            if abs(min(shiftMarkerList[mol][1]))<abs(max(shiftMarkerList[mol][1])):
                shiftY=max(shiftMarkerList[mol][1])
            shiftZ=min(shiftMarkerList[mol][2])
            if abs(min(shiftMarkerList[mol][2]))<abs(max(shiftMarkerList[mol][2])):
                shiftZ=max(shiftMarkerList[mol][2])
            shiftMarker[mol]=[shiftX,shiftY,shiftZ]
        #begin shift
        for mol in self.molList:
            for atom in self.molList[mol]:
                atom[3]=atom[3]-shiftMarker[mol][0]*cellX
                atom[4]=atom[4]-shiftMarker[mol][1]*cellY
                atom[5]=atom[5]-shiftMarker[mol][2]*cellZ
                
     #construct  atom cluster based on atom type and range      
        if self.groupType.upper()=='C':
            totalCluster=0
            for i in self.atomList:
                atom=map(float,i)
                clusterIndex=0
                for cluster in self.clusterList:
                    if atom in self.clusterList[cluster]:
                        clusterIndex=cluster
                if clusterIndex==0:
                    clusterIndex=totalCluster+1
                    self.clusterList[str(clusterIndex)]=[atom]    
                    totalCluster=clusterIndex   
                for j in self.atomList + self.atomListEx:
                    atom2=map(float,j)
                    atom2mX=atom2[:3]+[atom2[3]-cellX]+atom2[4:]
                    atom2pX=atom2[:3]+[atom2[3]+cellX]+atom2[4:]
                    atom2mY=atom2[:3]+[atom2[4]-cellY]+atom2[5:]  
                    atom2pY=atom2[:3]+[atom2[4]-cellY]+atom2[5:]
                    atom2mZ=atom2[:5]+[atom2[5]-cellZ]+atom2[6:]
                    atom2pZ=atom2[:5]+[atom2[5]-cellZ]+atom2[6:]
                    mirrorAtoms=[atom2,atom2mX,atom2pX,atom2mY,atom2pY,atom2mZ,atom2pZ]                
                    dist0=((atom[3]-atom2[3])**2+(atom[4]-atom2[4])**2+(atom[5]-atom2[5])**2)**0.5
                    dist1=((atom[3]-atom2mX[3])**2+(atom[4]-atom2[4])**2+(atom[5]-atom2[5])**2)**0.5
                    dist2=((atom[3]-atom2pX[3])**2+(atom[4]-atom2[4])**2+(atom[5]-atom2[5])**2)**0.5
                    dist3=((atom[3]-atom2[3])**2+(atom[4]-atom2mY[4])**2+(atom[5]-atom2[5])**2)**0.5
                    dist4=((atom[3]-atom2[3])**2+(atom[4]-atom2pY[4])**2+(atom[5]-atom2[5])**2)**0.5
                    dist5=((atom[3]-atom2[3])**2+(atom[4]-atom2[4])**2+(atom[5]-atom2mZ[5])**2)**0.5
                    dist6=((atom[3]-atom2[3])**2+(atom[4]-atom2[4])**2+(atom[5]-atom2pZ[5])**2)**0.5
                    minDistance, minIdx = min((val, idx) for (idx, val) in enumerate([dist0,dist1,dist2,dist3,dist4,dist5,dist6]))
                    if minDistance < self.clusterRange and minDistance > 0 and not atom2 in self.clusterList[str(clusterIndex)]:
                        self.clusterList[str(clusterIndex)].append(mirrorAtoms[minIdx]) 
            cluster2mol={}
            for cluster in self.clusterList:
                molId='-'.join(str(int(mol[0])) for mol in self.clusterList[cluster])
                cluster2mol[cluster]=cluster+'('+molId+')'
                #for mol in self.clusterList[cluster]:
                #    if str(int(mol[0]))=='1':
                #        print mol
            #print 'Cluster that identified as '+self.groupName+' are '+','.join([cluster2mol[cluster] for cluster in cluster2mol])+'. And the number of atoms in them are '+','.join([str(len(self.clusterList[cluster])) for cluster in self.clusterList]) 
                             
        #calculate mass center
        if self.groupType == 'M':
            for mol in self.molList:
                sumXM=0
                sumYM=0
                sumZM=0
                sumMass=0
                for atom in self.molList[mol]:
                    sumXM=sumXM+atom[2]*atom[3]
                    sumYM=sumYM+atom[2]*atom[4]
                    sumZM=sumZM+atom[2]*atom[5]
                    sumMass=sumMass+atom[2]
                xCenter=sumXM/sumMass
                yCenter=sumYM/sumMass
                zCenter=sumZM/sumMass
                self.massCenter[mol]=[xCenter,yCenter,zCenter]    
        if self.groupType == 'C':    
            for cluster in self.clusterList:
                sumXM=0
                sumYM=0
                sumZM=0
                sumMass=0
                for atom in self.clusterList[cluster]:
                    sumXM=sumXM+atom[2]*atom[3]
                    sumYM=sumYM+atom[2]*atom[4]
                    sumZM=sumZM+atom[2]*atom[5]
                    sumMass=sumMass+atom[2]
                xCenter=sumXM/sumMass
                yCenter=sumYM/sumMass
                zCenter=sumZM/sumMass
                self.massCenter[cluster]=[xCenter,yCenter,zCenter]  
    def printInfo(self):
        if self.groupType.upper()=='C':
            cluster2mol={}
            for cluster in self.clusterList:
                molId='-'.join(str(int(mol[0])) for mol in self.clusterList[cluster])
                cluster2mol[cluster]=cluster+'('+molId+')'
            print str(len(cluster2mol))+' Clusters that identified as '+self.groupName+' are '+','.join([cluster2mol[cluster] for cluster in cluster2mol])+'. And the number of atoms in them are '+','.join([str(len(self.clusterList[cluster])) for cluster in self.clusterList]) 
        if self.groupType.upper()=='M':
            print str(len(self.molList))+' Mol groups that identified as '+self.groupName+' are '+','.join([mol for mol in self.molList])+'. And the number of atoms in them are '+','.join([str(len(self.molList[mol])) for mol in self.molList])     
            
            
    def plot3d(self):
        #fig = plt.figure('1')
        #plt.subplot(111, projection='3d')
        xCenter=[]
        yCenter=[]
        zCenter=[]
        for mol in self.massCenter:
            xCenter.append(self.massCenter[mol][0])
            yCenter.append(self.massCenter[mol][1])
            zCenter.append(self.massCenter[mol][2])
        return  xCenter,yCenter,zCenter
           

#fig=plt.figure()
#ax=fig.add_subplot(111, projection='3d')
#massG=massCenterGroup('D:M:3')
#massG.extractInfo(trjData[0],indexList,characterLen)
#xCenter,yCenter,zCenter=massG.plot3d()


#ax.scatter(xCenter,yCenter,zCenter,c=lineColors[1],marker='o',label=massG.groupName,)


class AnimatedScatter(object):
    def __init__(self,inputFile,dataFile):
        self.inputFile=inputFile
        self.group=[]
        self.measure=[]
        self.frameRange=[]
        groupSwitch=0
        measureSwitch=0
        for line in open(inputFile,'r'):
            if '#group' in line:
                groupSwitch=1
                measureSwitch=0
                continue
            elif '#measure' in line:
                groupSwitch=0
                measureSwitch=1
                continue
            elif '#' in line:
                continue
            if groupSwitch==1:
                self.group.append(line)
            if measureSwitch==1:
                self.measure.append(line)
        self.angle = 0
        self.rawData=dataFile
        self.fig = plt.figure()
        self.ax = self.fig.add_subplot(111,projection = '3d')
        #read the trajectory data
        self.trjData=[]
        self.trjFile=open(self.rawData,'r')
        blockID=-1
        for line in self.trjFile:
            if 'ITEM: TIMESTEP' in line:
                self.trjData=self.trjData+[[]]
                blockID=blockID+1
            self.trjData[blockID].append(line)
        self.totalFrame=len(self.trjData)
        #read the head line in trjData and generate indexlist and characterlen
        self.indexList=[]
        keywords=['mol','type','mass','xu','yu','zu','ix','iy','iz']
        for line in self.trjData[0]:
            if 'ITEM: ATOMS' in line:
                self.headerList=line[12:].split()
        try:
            for i in keywords:
                self.indexList.append(self.headerList.index(i))
        except ValueError as e:
            print e[0].split()[0]+' is not found in the column header: '+','.join(self.headerList) 
            sys.exit()
        self.characterLen=len(self.headerList)    
              
        #generate color list for different group
        step=sum(1 for line in self.group)
        rg=step/2+1
        gb=step/2+1
        if rg+gb>step+1:
            gb=gb-1
        rglist=[]
        gblist=[]
        for i in range(rg):
            rglist.append((1-i/(rg-1.0),i/(rg-1.0),0.0))
        for i in range(gb):
            gblist.append((0.0,1.0-i/(rg-1.0),i/(rg-1.0)))
        rglist.pop()
        self.rgblist=rglist+gblist           
            
    def genAni(self,interval): 
        self.ani = animation.FuncAnimation(self.fig, self.animate, interval=interval,init_func=self.setup_plot, blit=True)
        self.ax.legend()
        plt.show()
        #self.ani.save('demoanimation.gif', writer='imagemagick', fps=4);
    def setup_plot(self):
        self.scat={}
        j=0
        for line in self.group:    
            print 'INPUT: '+line.replace('\n','')
            massG=massCenterGroup(line.replace('\n',''))
            massG.extractInfo(self.trjData[0],self.indexList,self.characterLen)
            xCenter,yCenter,zCenter=massG.plot3d()
            self.scat[massG.groupName]=self.ax.scatter(xCenter,yCenter,zCenter,c=self.rgblist[j],marker='o',s=80,label=massG.groupName,animated=True)
            
            j=j+1
        
        return self.scat.values()
    
    def plotFrame(self,frame=0):
        self.scat={}
        j=0
        print 'Frame: '+ str(frame)
        for line in self.group:    
            
            massG=massCenterGroup(line.replace('\n',''))
            massG.extractInfo(self.trjData[frame],self.indexList,self.characterLen)
            xCenter,yCenter,zCenter=massG.plot3d()
            self.scat[massG.groupName]=self.ax.scatter(xCenter,yCenter,zCenter,c=self.rgblist[j],marker='o',s=80,label=massG.groupName)
            j=j+1
        self.ax.legend()  
        plt.show()
             
                    
                        
                                
    def animate(self,i):
        j=0
        print "Frame: "+str(i)
        for line in self.group:
            #print 'INPUT: '+line.replace('\n','')
            massG=massCenterGroup(line.replace('\n',''))
            massG.extractInfo(self.trjData[i],self.indexList,self.characterLen)
            xCenter,yCenter,zCenter=massG.plot3d()
            self.scat[massG.groupName]._offsets3d = (xCenter,yCenter,zCenter)
            plt.draw()
            j=j+1
        return self.scat.values()
     
    def checkGroup(self,frames=[0]):
        currentGroup={}
        currentAtom={}
        lastGroup={}
        lastAtom={}
        threshold=0.02
        for line in self.group:
            currentGroup[line.split(':')[0]]=[]
            currentAtom[line.split(':')[0]]=[]
            lastGroup[line.split(':')[0]]=[]
            lastAtom[line.split(':')[0]]=[]
        for i in frames:
            if float(i)/len(frames) > threshold:
                print 'Checking frame ' + str(i)
                threshold=threshold+0.02 
            for line in self.group:
                massG=massCenterGroup(line.replace('\n',''))
                massG.extractInfo(self.trjData[i],self.indexList,self.characterLen)
                if massG.groupType.upper()=='M':
                    currentGroup[line.split(':')[0]]=[mol for mol in massG.molList]
                    currentAtom[line.split(':')[0]]=[str(len(massG.molList[mol])) for mol in massG.molList]
                    if not lastGroup[line.split(':')[0]]:
                        lastGroup[line.split(':')[0]]=currentGroup[line.split(':')[0]]
                        lastAtom[line.split(':')[0]]=currentAtom[line.split(':')[0]]
                    elif lastGroup[line.split(':')[0]] != currentGroup[line.split(':')[0]] or lastAtom[line.split(':')[0]] != currentAtom[line.split(':')[0]]:
                        print "Mol group "+massG.groupName+" in frame "+str(i)+" is not equal to atom group in frame "+str(i-1)+", exit now"
                        sys.exit()
                    #print str(len(massG.molList))+' Molecules that identified as '+massG.groupName+' are '+','.join([mol for mol in massG.molList])+'. And the number of atoms in them are '+','.join([str(len(massG.molList[mol])) for mol in massG.molList])             
                if massG.groupType.upper()=='C':
                    currentGroup[line.split(':')[0]]=[cluster for cluster in massG.clusterList]
                    currentAtom[line.split(':')[0]]=[[str(len(massG.clusterList[cluster])) for cluster in massG.clusterList]]
                    if not lastGroup[line.split(':')[0]]:
                        lastGroup[line.split(':')[0]]=currentGroup[line.split(':')[0]]
                        lastAtom[line.split(':')[0]]=currentAtom[line.split(':')[0]]
                    elif lastGroup[line.split(':')[0]] != currentGroup[line.split(':')[0]] or lastAtom[line.split(':')[0]] != currentAtom[line.split(':')[0]]:
                        print "Cluster group "+massG.groupName+" in frame "+str(i)+" is not equal to atom group in frame "+str(i-1)+", exit now"
                        sys.exit()
        for line in self.group:
            massG=massCenterGroup(line.replace('\n',''))
            massG.extractInfo(self.trjData[frames[0]],self.indexList,self.characterLen)
            massG.printInfo()
        sys.exit()
    
    def massCenterDist(self,group1,group2):
        nearCenterSum={}
        nearCenterStat={}
        nearCenterCount={}
        allDistStat={}
        allDistList={}
        threshold=0.05
        if not self.frameRange:
            self.frameRange=[self.totalFrame]
        for i in range(*self.frameRange):          
            if float(i)/len(range(*self.frameRange)) > threshold:
                print 'Processing frame ' + str(i)
                threshold=threshold+0.05 
            groups={}
            for line in self.group:
                massG=massCenterGroup(line.replace('\n',''))
                massG.extractInfo(self.trjData[i],self.indexList,self.characterLen)
                groups[massG.groupName]=massG
            cellX=massG.cellX
            cellY=massG.cellY
            cellZ=massG.cellZ
            firstGroup=group1.split('+')
            secondGroup=group2.split('+')
            allDist={}
            nearCenterList={}
            for groupA in firstGroup:
                for mc1 in groups[groupA].massCenter:
                    #a dictionary, the key is the every center in groupA and the value is another dict which contains the distance to nearest center in groupB
                    allDist[groupA+'_'+mc1]={}
                    nearCenterList[groupA+'_'+mc1]={}
                    #a dictionary, the key is  in groupB and the value is the count 
                    minDistOnly={}
                    minDist={}
                    for groupB in secondGroup:
                        #a list contain all the distance from a center in groupA to every center in groupB
                        dist2groupB=[]
                        #a list contain all the mass center in group B
                        mc2serial=[]
                        for mc2 in groups[groupB].massCenter:
                            c1=groups[groupA].massCenter[mc1]
                            c2=groups[groupB].massCenter[mc2]
                            mc2serial.append(mc2)
                            c1mX=[c1[0]-cellX]+c1[1:]
                            c1pX=[c1[0]+cellX]+c1[1:]
                            c1mY=c1[:1]+[c1[1]-cellY]+c1[2:]
                            c1pY=c1[:1]+[c1[1]+cellY]+c1[2:] 
                            c1mZ=c1[:2]+[c1[2]-cellZ]
                            c1pZ=c1[:2]+[c1[2]+cellZ]             
                            dist0=((c1[0]-c2[0])**2+(c1[1]-c2[1])**2+(c1[2]-c2[2])**2)**0.5
                            dist1=((c1mX[0]-c2[0])**2+(c1[1]-c2[1])**2+(c1[2]-c2[2])**2)**0.5
                            dist2=((c1pX[0]-c2[0])**2+(c1[1]-c2[1])**2+(c1[2]-c2[2])**2)**0.5
                            dist3=((c1[0]-c2[0])**2+(c1mY[1]-c2[1])**2+(c1[2]-c2[2])**2)**0.5
                            dist4=((c1[0]-c2[0])**2+(c1pY[1]-c2[1])**2+(c1[2]-c2[2])**2)**0.5
                            dist5=((c1[0]-c2[0])**2+(c1[1]-c2[1])**2+(c1mZ[2]-c2[2])**2)**0.5
                            dist6=((c1[0]-c2[0])**2+(c1[1]-c2[1])**2+(c1pZ[2]-c2[2])**2)**0.5
                            minDistance=min([dist0,dist1,dist2,dist3,dist4,dist5,dist6])
                            dist2groupB.append(minDistance)
                        minDistance, minIdx = min((val, idx) for (idx, val) in enumerate(dist2groupB))
                        minDist[groupB]=[minDistance,mc2serial[minIdx]]
                    for group in minDist:
                        minDistOnly[group]=[]
                    nearGroup=min(minDist.items(), key=lambda x: x[1][0])
                    minDistOnly[nearGroup[0]]=[nearGroup[1][0]]
                    allDist[groupA+'_'+mc1]=minDist
                    nearCenterList[groupA+'_'+mc1]=minDistOnly
                
            #put all the distance into a list        
            for center1 in allDist:
                if center1 not in allDistList:
                    allDistList[center1]={}
                for center2 in allDist[center1]:
                    try:
                        allDistList[center1][center2].append(allDist[center1][center2][0])
                    except KeyError:
                        allDistList[center1][center2]=[allDist[center1][center2][0]]
            

            
            for center1 in nearCenterList:
                if center1 not in nearCenterSum:
                    nearCenterSum[center1]={}
                for center2 in nearCenterList[center1]:
                    try:
                        if nearCenterList[center1][center2]:
                            nearCenterSum[center1][center2]=nearCenterSum[center1][center2]+nearCenterList[center1][center2]
                    except KeyError:
                        nearCenterSum[center1][center2]=nearCenterList[center1][center2]
         
        for center1 in nearCenterSum:
            if center1 not in nearCenterStat:
                nearCenterStat[center1]={}
            for center2 in nearCenterSum[center1]:
                nearCenterStat[center1][center2]=[np.mean(np.array(nearCenterSum[center1][center2])),np.std(np.array(nearCenterSum[center1][center2]))]
        
        for center1 in nearCenterSum:
            if center1 not in nearCenterCount:
                nearCenterCount[center1]={}
            for center2 in nearCenterSum[center1]:
                nearCenterCount[center1][center2]=len(nearCenterSum[center1][center2])
        
        for center1 in allDistList:
                if center1 not in allDistStat:
                    allDistStat[center1]={}
                for center2 in allDistList[center1]:
                    allDistStat[center1][center2]=[np.mean(np.array(allDistList[center1][center2])),np.std(np.array(allDistList[center1][center2]))]
       
        print "\n====================Nearest Center STAT (nearest group)===================="
        for i in range(len(nearCenterStat)):
            if i==0:
                print ' '.join(['{:<20}']*(len(nearCenterStat.values()[0])+1)).format(" ",*[j for j in nearCenterStat.values()[0]])
            print ('{:<20s}'+' '.join(['{:<20s}']*len(nearCenterStat.values()[0]))).format(nearCenterStat.keys()[i]+'(mean/sd)',*['{:.6f}'.format(nearCenterStat.values()[i][j][0])+'/'+'{:.6f}'.format(nearCenterStat.values()[i][j][1]) for j in nearCenterStat.values()[i]])
        print "\n====================Nearest Center Count===================="   
        for i in range(len(nearCenterCount)):
            if i==0:
                print ' '.join(['{:<20}']*(len(nearCenterCount.values()[0])+1)).format(" ",*[j for j in nearCenterCount.values()[0]])
        print ('{:<20s}'+' '.join(['{:<20s}']*len(nearCenterCount.values()[0]))).format(nearCenterCount.keys()[i]+'(count)',*['{:d}'.format(nearCenterCount.values()[i][j]) for j in nearCenterCount.values()[i]])    
        print "\n====================Nearest Center STAT (every group)===================="       
        for i in range(len(allDistStat)):
            if i==0:
                print ' '.join(['{:<20}']*(len(allDistStat.values()[0])+1)).format(" ",*[j for j in allDistStat.values()[0]])
        print ('{:<20s}'+' '.join(['{:<20s}']*len(allDistStat.values()[0]))).format(allDistStat.keys()[i]+'(mean/sd)',*['{:.6f}'.format(allDistStat.values()[i][j][0])+'/'+'{:.6f}'.format(allDistStat.values()[i][j][1]) for j in allDistStat.values()[i]])
        
        
        
        #print nearCenterStat
        #print nearCenterCount
        #print allDistStat
        #
                    
            
parser=argparse.ArgumentParser(description='analysis trajectory data')
parser.add_argument("-i","--input",help='The name of input analysis command for trajectory',type=str)
parser.add_argument("-c","--checkGroup",help='check the defined group, input numbers for frame, e.g. -c 1 check frame 1; -c 5,100 check frame 5 to 100; -c all check all frames',type=str)
parser.add_argument("-a","--animation",help='start the animation with the input inteval, e.g. -a 1000 the animation has an interval of 1 second',type=str)
parser.add_argument("-p","--plot",help='plot a specific frame, e.g. -p 100 plot frame 100',type=int)
parser.add_argument("-f","--frames",help='analysis only designate frame, e.g. -f 1,100 analysis frame 1 to frame 100',type=str)
parser.add_argument("trajectory",nargs='+')
args=parser.parse_args()
                
stat = AnimatedScatter(args.input,args.trajectory[0])

if args.checkGroup:
    if args.checkGroup=='all':
        frames=range(stat.totalFrame)
    else:
        frames=range(*map(int,args.checkGroup.split(',')))
    stat.checkGroup(frames)
    
if args.animation:
    stat.genAni(args.animation)
    sys.exit()

if args.plot:
    stat.plotFrame(args.plot)
    sys.exit()
if args.frames:
    stat.frameRange=map(int,args.frames.split(','))
print "Total frame number is " + str(stat.totalFrame)

for line in stat.measure:
    if line.rstrip().split(':')[0]=='mcdist':        
        stat.massCenterDist(line.rstrip().split(':')[1],line.rstrip().split(':')[2])
 
   
